import torch
import torch.nn as nn
import torch.nn.functional as F
from utils import build_targets
from cfg import yolo_cfg

class Upsample(nn.Module):
    def __init__(self, scale_factor, mode="nearest"):
        super(Upsample, self).__init__()
        self.scale_factor = scale_factor
        self.mode = mode

    def forward(self, x):
        x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
        return x


class DBL(nn.Module):
    def __init__(self,in_chanel, out_chanel, kernel_size=(3,3), padding=1, stride=1):
        super(DBL,self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_chanel, out_chanel, kernel_size, stride=stride, padding=padding),
            nn.BatchNorm2d(out_chanel),
            nn.LeakyReLU()
        )

    def forward(self,x):
        return self.model(x)


class Res_unit(nn.Module):
    def __init__(self, chanel, stride):
        super(Res_unit, self).__init__()
        self.DBL_1 = DBL(chanel, chanel//2, kernel_size=(1,1), padding=0, stride=stride)
        self.DBL_2 = DBL(chanel//2, chanel, kernel_size=(3,3),padding=1, stride=stride)

    def forward(self,x):
        x_ = x.clone()
        x = self.DBL_2(self.DBL_1(x)) + x_
        return x

class Resblock_body(nn.Module):
    def __init__(self, units, in_chanel, out_chanel, stride=1):
        super(Resblock_body, self).__init__()
        self.model = nn.Sequential(
            DBL(in_chanel=in_chanel, out_chanel=out_chanel, stride=2)
        )
        for i in range(units):
            self.model.add_module("res_1", Res_unit(out_chanel, stride))
    def forward(self,x):
        return self.model(x)


class Yolo_head(nn.Module):
    def __init__(self, anchors, grid_size, num_classes=80, image_dim=416):
        super(Yolo_head, self).__init__()
        self.anchors = anchors
        self.grid_size = grid_size
        self.num_classes = num_classes
        self.img_dim = image_dim
        self.num_anchors = len(anchors)
        self.anchor_prepred = False

        # loss function
        self.mse_loss = nn.MSELoss()
        self.bce_loss = nn.BCELoss()

        # 一些超参
        self.obj_scale = 1 # conf正负比重
        self.noobj_scale = 100

    def prepare_anchor(self, x):
        self.anchor_prepred = True
        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor

        g = x.size(-1)
        self.stride = self.img_dim / self.grid_size
        self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor)
        self.grid_y = torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor)
        self.scaled_anchors = FloatTensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors])
        self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))
        self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))

    def forward(self, x, targets=None):

        # Tensors for cuda support

        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor

        # 长宽
        grid_size = x.size(-1)
        num_samples = x.size(0)

        prediction = (
            x.view(num_samples, self.num_anchors, self.num_classes + 5, grid_size, grid_size)
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )

        # Get outputs
        x = torch.sigmoid(prediction[..., 0])  # Center x
        y = torch.sigmoid(prediction[..., 1])  # Center y
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height
        pred_conf = torch.sigmoid(prediction[..., 4])  # Conf
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

        if not self.anchor_prepred:
            self.prepare_anchor(x)

        # Add offset and scale with anchors
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + self.grid_x
        pred_boxes[..., 1] = y.data + self.grid_y
        pred_boxes[..., 2] = torch.exp(w.data) * self.anchor_w
        pred_boxes[..., 3] = torch.exp(h.data) * self.anchor_h

        output = torch.cat(
            (
                pred_boxes.view(num_samples, -1, 4) * self.stride,
                pred_conf.view(num_samples, -1, 1),
                pred_cls.view(num_samples, -1, self.num_classes),
            ),
            -1,
        )
        if targets is None:
            return output, 0
        else:
            # TODO(把这里的代码重新实现)
            iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf = build_targets(
                pred_boxes=pred_boxes,
                pred_cls=pred_cls,
                target=targets,
                anchors=self.scaled_anchors,
                ignore_thres=0.5,
            )
            # Loss : Mask outputs to ignore non-existing objects (except with conf. loss)
            loss_x = self.mse_loss(x[obj_mask], tx[obj_mask])
            loss_y = self.mse_loss(y[obj_mask], ty[obj_mask])
            loss_w = self.mse_loss(w[obj_mask], tw[obj_mask])
            loss_h = self.mse_loss(h[obj_mask], th[obj_mask])
            loss_conf_obj = self.bce_loss(pred_conf[obj_mask], tconf[obj_mask])
            loss_conf_noobj = self.bce_loss(pred_conf[noobj_mask], tconf[noobj_mask])
            loss_conf = self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj
            loss_cls = self.bce_loss(pred_cls[obj_mask], tcls[obj_mask])
            total_loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
            # print(
            #     "loss_x="+str(loss_x.data)+"\n",
            #     "loss_y=" + str(loss_y.data)+"\n",
            #     "loss_w=" + str(loss_w.data)+"\n",
            #     "loss_h=" + str(loss_h.data)+"\n",
            #     "loss_conf=" + str(loss_conf.data)+"\n",
            #     "loss_cls=" + str(loss_cls.data)+"\n",
            #     )
            return output, total_loss

class Darknet(nn.Module):
    def __init__(self, in_chanel, use_fc=False, fc_out=10):
        super(Darknet, self).__init__()
        self.DBL = DBL(3, 32)
        self.res_123 = nn.Sequential(
            Resblock_body(1, 32, 64),
            Resblock_body(2, 64, 128),
            Resblock_body(8, 128, 256)
        )
        self.res4_8x = Resblock_body(8, 256, 512)
        self.res5_4x = Resblock_body(4, 512, 1024)
        self.pool = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Linear(1024, fc_out)
        self.use_fc = use_fc

        # yolo 多尺度融合层及输出层
        # darknet最后一个block直接经过一个DBL_block和DBL_conv_1输出y1
        self.DBL_block_1 = nn.Sequential(
            DBL(1024, 512, (1,1), padding=0, stride=1),
            DBL(512, 1024, (3, 3), padding=1, stride=1),
            DBL(1024, 512, (1, 1), padding=0, stride=1),
            DBL(512, 1024, (3, 3), padding=1, stride=1),
            DBL(1024, 512, (1, 1), padding=0, stride=1),
        )
        self.DBL_conv_1 = nn.Sequential(
            DBL(512, 1024, (3,3), padding=1, stride=1),
            nn.Conv2d(1024, 255, (1,1), padding=0, stride=1)
        )

        # darknet倒数第二个block融合最后一个block的结果后，经过一个DBL_block和DBL_conv_1输出y2
        self.DBL_upsample_2 = nn.Sequential(
            DBL(512, 256, (1,1), padding=0, stride=1),
            Upsample(2),
        )
        self.DBL_block_2 = nn.Sequential(
            DBL(512+256, 256, (1, 1), padding=0, stride=1),
            DBL(256, 512, (3, 3), padding=1, stride=1),
            DBL(512, 256, (1, 1), padding=0, stride=1),
            DBL(256, 512, (3, 3), padding=1, stride=1),
            DBL(512, 256, (1, 1), padding=0, stride=1),
        )
        self.DBL_conv_2 = nn.Sequential(
            DBL(256, 512, (3,3), padding=1, stride=1),
            nn.Conv2d(512, 255, (1,1), padding=0, stride=1)
        )

        # darknet倒数第三个block融合上一个block的结果后，经过一个DBL_block和DBL_conv_1输出y3
        self.DBL_upsample_3 = nn.Sequential(
            DBL(256, 128, (1,1), padding=0, stride=1),
            Upsample(2),
        )
        self.DBL_block_3 = nn.Sequential(
            DBL(256+128, 128, (1, 1), padding=0, stride=1),
            DBL(128, 256, (3, 3), padding=1, stride=1),
            DBL(256, 128, (1, 1), padding=0, stride=1),
            DBL(128, 256, (3, 3), padding=1, stride=1),
            DBL(256, 128, (1, 1), padding=0, stride=1),
        )
        self.DBL_conv_3 = nn.Sequential(
            DBL(128, 256, (3,3), padding=1, stride=1),
            nn.Conv2d(256, 255, (1,1), padding=0, stride=1)
        )
        # Yolo head
        # TODO(改成读取参数中的anchor)
        self.yolo_heads = [Yolo_head(anchors=yolo_cfg["anchors"][6:], grid_size=13, ),
                           Yolo_head(anchors=yolo_cfg["anchors"][3:6], grid_size=26, ),
                           Yolo_head(anchors=yolo_cfg["anchors"][0:3], grid_size=52, )]

    def forward(self,x, targets=None):
        x = self.DBL(x)
        x = x3_256 = self.res_123(x)
        x = x4_512 = self.res4_8x(x)
        x = x5_1024 = self.res5_4x(x)

        # 仅预训练使用
        if self.use_fc:
            x = self.fc(self.pool(x).view(x.shape[0],-1))
            x = F.log_softmax(x)
            return x
        else:
            # 多尺度融合需要从x5_1024开始计算
            x5_512 = self.DBL_block_1(x5_1024)
            y1 = self.DBL_conv_1(x5_512)

            x4_256 = self.DBL_block_2(
                torch.cat(
                    (self.DBL_upsample_2(x5_512), x4_512,),
                    1
                )
            )
            y2 = self.DBL_conv_2(x4_256)

            x3_128 = self.DBL_block_3(
                torch.cat(
                    (self.DBL_upsample_3(x4_256), x3_256,),
                    1
                )
            )
            y3 = self.DBL_conv_3(x3_128)

            y1, loss1 = self.yolo_heads[0](y1, targets)
            y2, loss2 = self.yolo_heads[1](y2, targets)
            y3, loss3 = self.yolo_heads[2](y3, targets)
            #
            # y1 = torch.cat((y1,torch.Tensor([1 for i in range(y1.shape[1])]).unsqueeze(1).repeat(y1.shape[0],1,1)),dim=2)
            # y2 = torch.cat((y2,torch.Tensor([2 for i in range(y2.shape[1])]).unsqueeze(1).repeat(y2.shape[0],1,1)),dim=2)
            # y3 = torch.cat((y3,torch.Tensor([3 for i in range(y3.shape[1])]).unsqueeze(1).repeat(y3.shape[0],1,1)),dim=2)

            return torch.cat((y1,y2,y3),1), loss1+loss2+loss3